"""
python fujian1_pmdarima_afterSplineAndSmooth.py
"""
import os
import json
import pandas as pd
from pmdarima import auto_arima
from tqdm import tqdm
import logging

# 设置输入和输出目录（使用相对路径）
input_dir = r"./fujian/fujian1/spline_then_smooth"
output_dir = r"./fujian/fujian1/pmdarima"
logs_dir = os.path.join(output_dir, 'logs')

# 创建输出和日志目录（如果不存在）
os.makedirs(output_dir, exist_ok=True)
os.makedirs(logs_dir, exist_ok=True)

# 设置日志记录，仅在有警告或错误时记录
error_log_path = os.path.join(logs_dir, 'errs', 'all_err.log')
os.makedirs(os.path.dirname(error_log_path), exist_ok=True)

# 日志配置，仅在发生警告或错误时记录
logging.basicConfig(filename=os.path.join(logs_dir, 'category.log'), level=logging.WARNING,
                    format='%(asctime)s - %(levelname)s - %(message)s')

# 遍历目录中的所有 JSON 文件
json_files = [f for f in os.listdir(input_dir) if f.endswith('.json')]
total_files = len(json_files)

for idx, filename in tqdm(enumerate(json_files), total=total_files, desc="Processing files"):
    file_path = os.path.join(input_dir, filename)

    try:
        # 读取 JSON 文件
        with open(file_path, 'r') as f:
            data = json.load(f)

        # 将数据转换为 DataFrame
        df = pd.DataFrame(data)
        df['date'] = pd.to_datetime(df['date'])  # 转换日期格式
        df.set_index('date', inplace=True)

        # 确保有 12 个数据点进行训练
        if len(df) != 12:
            logging.warning(f"Not enough data points in {filename} for ARIMA model.")
            continue

        # 填充缺失值（如果有）
        df['inventory'] = df['inventory'].ffill()  # 向前填充缺失值

        # 检查是否存在 NaN 或零值
        if df['inventory'].isnull().any():
            logging.error(f"Data contains NaN or zero in {filename}. Skipping this file.")
            with open(error_log_path, 'a') as error_log:
                error_log.write(f"{filename}\n")
            continue

        # 使用 auto_arima 进行模型选择
        model = auto_arima(df['inventory'], seasonal=False, stepwise=True, suppress_warnings=True)

        # 进行预测，预测未来 3 个数据点
        forecast, conf_int = model.predict(n_periods=3, return_conf_int=True)

        # 生成未来日期
        forecast_dates = pd.date_range(start=df.index[-1] + pd.DateOffset(months=1), periods=3, freq='MS')

        # 创建预测结果的 DataFrame
        forecast_result = pd.DataFrame({'date': forecast_dates, 'inventory': forecast})

        # 合并原始数据和预测数据
        result_df = pd.concat([df.reset_index(), forecast_result], ignore_index=True)

        # 转换日期为字符串格式以便 JSON 序列化
        result_df['date'] = result_df['date'].dt.strftime('%Y-%m-%d')

        # 转换回 JSON 格式并输出
        output_path = os.path.join(output_dir, filename)
        result_json = result_df.to_dict(orient='records')

        with open(output_path, 'w') as f:
            json.dump(result_json, f, indent=4)

    except Exception as e:
        logging.error(f"Error processing {filename}: {e}")
        with open(error_log_path, 'a') as error_log:
            error_log.write(f"文件名：{filename}，异常输出： {str(e)}\n")
